package com.rrd.hive.udf.brickhouse.udf.collect;

import java.util.ArrayList;
import java.util.HashMap;

import org.apache.hadoop.hive.ql.exec.Description;
import org.apache.hadoop.hive.ql.exec.UDFArgumentTypeException;
import org.apache.hadoop.hive.ql.metadata.HiveException;
import org.apache.hadoop.hive.ql.parse.SemanticException;
import org.apache.hadoop.hive.ql.udf.generic.AbstractGenericUDAFResolver;
import org.apache.hadoop.hive.ql.udf.generic.GenericUDAFEvaluator;
import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspectorFactory;
import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspectorUtils;
import org.apache.hadoop.hive.serde2.objectinspector.PrimitiveObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.StandardListObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.StandardMapObjectInspector;
import org.apache.hadoop.hive.serde2.typeinfo.TypeInfo;

@Description(name = "collect",
        value = "_FUNC_(x) - Returns an array of all the elements in the aggregation group ")
public class CollectUDAF extends AbstractGenericUDAFResolver {

    @Override
    public GenericUDAFEvaluator getEvaluator(TypeInfo[] parameters)
            throws SemanticException {
        // TODO Auto-generated method stub
        if (parameters.length != 1 && parameters.length != 2) {
            throw new UDFArgumentTypeException(parameters.length - 1,
                "One argument is expected to return an Array, Two arguments are expected for a Map.");
        }
        if (parameters.length == 1) {
            return new ArrayCollectUDAFEvaluator();
        } else {
            return new MapCollectUDAFEvaluator();
        }
    }

    public static class ArrayCollectUDAFEvaluator extends GenericUDAFEvaluator {
        // For PARTIAL1 and COMPLETE: ObjectInspectors for original data
        private ObjectInspector inputOI;
        // For PARTIAL2 and FINAL: ObjectInspectors for partial aggregations (list
        // of objs)
        private StandardListObjectInspector loi;
        private StandardListObjectInspector internalMergeOI;

        static class ArrayAggBuffer implements AggregationBuffer {
            ArrayList collectArray = new ArrayList();
        }

        @Override
        public ObjectInspector init(Mode m, ObjectInspector[] parameters)
                throws HiveException {
            super.init(m, parameters);
            // init output object inspectors
            // The output of a partial aggregation is a list
            if (m == Mode.PARTIAL1) {
                this.inputOI = parameters[0];
                return ObjectInspectorFactory
                    .getStandardListObjectInspector(ObjectInspectorUtils
                        .getStandardObjectInspector(this.inputOI));
            } else {
                if (!(parameters[0] instanceof StandardListObjectInspector)) {
                    //no map aggregation.
                    this.inputOI = ObjectInspectorUtils
                        .getStandardObjectInspector(parameters[0]);
                    return ObjectInspectorFactory
                        .getStandardListObjectInspector(this.inputOI);
                } else {
                    this.internalMergeOI = (StandardListObjectInspector) parameters[0];
                    this.inputOI = this.internalMergeOI
                        .getListElementObjectInspector();
                    this.loi = (StandardListObjectInspector) ObjectInspectorUtils
                        .getStandardObjectInspector(this.internalMergeOI);
                    return this.loi;
                }
            }
        }

        @Override
        public AggregationBuffer getNewAggregationBuffer()
                throws HiveException {
            AggregationBuffer buff = new ArrayAggBuffer();
            this.reset(buff);
            return buff;
        }

        @Override
        public void iterate(AggregationBuffer agg, Object[] parameters)
                throws HiveException {
            Object p = parameters[0];

            if (p != null) {
                ArrayAggBuffer myagg = (ArrayAggBuffer) agg;
                this.putIntoSet(p, myagg);
            }
        }

        @Override
        public void merge(AggregationBuffer agg, Object partial)
                throws HiveException {
            ArrayAggBuffer myagg = (ArrayAggBuffer) agg;
            ArrayList<Object> partialResult = (ArrayList<Object>) this.internalMergeOI
                .getList(partial);
            for (Object i : partialResult) {
                this.putIntoSet(i, myagg);
            }
        }

        @Override
        public void reset(AggregationBuffer buff) throws HiveException {
            ArrayAggBuffer arrayBuff = (ArrayAggBuffer) buff;
            arrayBuff.collectArray = new ArrayList();
        }

        @Override
        public Object terminate(AggregationBuffer agg) throws HiveException {
            ArrayAggBuffer myagg = (ArrayAggBuffer) agg;
            ArrayList<Object> ret = new ArrayList<>(myagg.collectArray.size());
            ret.addAll(myagg.collectArray);
            return ret;

        }

        private void putIntoSet(Object p, ArrayAggBuffer myagg) {
            Object pCopy = ObjectInspectorUtils.copyToStandardObject(p,
                this.inputOI);
            myagg.collectArray.add(pCopy);
        }

        @Override
        public Object terminatePartial(AggregationBuffer agg)
                throws HiveException {
            ArrayAggBuffer myagg = (ArrayAggBuffer) agg;
            ArrayList<Object> ret = new ArrayList<>(myagg.collectArray.size());
            ret.addAll(myagg.collectArray);
            return ret;
        }
    }

    public static class MapCollectUDAFEvaluator extends GenericUDAFEvaluator {
        // For PARTIAL1 and COMPLETE: ObjectInspectors for original data
        private PrimitiveObjectInspector inputKeyOI;
        private ObjectInspector inputValOI;
        // For PARTIAL2 and FINAL: ObjectInspectors for partial aggregations (list
        // of objs)
        private StandardMapObjectInspector moi;
        private StandardMapObjectInspector internalMergeOI;

        static class MapAggBuffer implements AggregationBuffer {
            HashMap<Object, Object> collectMap = new HashMap<>();
        }

        @Override
        public ObjectInspector init(Mode m, ObjectInspector[] parameters)
                throws HiveException {
            super.init(m, parameters);
            // init output object inspectors
            // The output of a partial aggregation is a list
            if (m == Mode.PARTIAL1) {
                this.inputKeyOI = (PrimitiveObjectInspector) parameters[0];
                this.inputValOI = parameters[1];

                return ObjectInspectorFactory.getStandardMapObjectInspector(
                    ObjectInspectorUtils
                        .getStandardObjectInspector(this.inputKeyOI),
                    ObjectInspectorUtils
                        .getStandardObjectInspector(this.inputValOI));
            } else {
                if (!(parameters[0] instanceof StandardMapObjectInspector)) {
                    this.inputKeyOI = (PrimitiveObjectInspector) ObjectInspectorUtils
                        .getStandardObjectInspector(parameters[0]);
                    this.inputValOI = ObjectInspectorUtils
                        .getStandardObjectInspector(parameters[1]);
                    return ObjectInspectorFactory.getStandardMapObjectInspector(
                        this.inputKeyOI, this.inputValOI);
                } else {
                    this.internalMergeOI = (StandardMapObjectInspector) parameters[0];
                    this.inputKeyOI = (PrimitiveObjectInspector) this.internalMergeOI
                        .getMapKeyObjectInspector();
                    this.inputValOI = this.internalMergeOI
                        .getMapValueObjectInspector();
                    this.moi = (StandardMapObjectInspector) ObjectInspectorUtils
                        .getStandardObjectInspector(this.internalMergeOI);
                    return this.moi;
                }
            }
        }

        @Override
        public AggregationBuffer getNewAggregationBuffer()
                throws HiveException {
            AggregationBuffer buff = new MapAggBuffer();
            this.reset(buff);
            return buff;
        }

        @Override
        public void iterate(AggregationBuffer agg, Object[] parameters)
                throws HiveException {
            Object k = parameters[0];
            Object v = parameters[1];

            if (k != null) {
                MapAggBuffer myagg = (MapAggBuffer) agg;
                this.putIntoSet(k, v, myagg);
            }
        }

        @Override
        public void merge(AggregationBuffer agg, Object partial)
                throws HiveException {
            MapAggBuffer myagg = (MapAggBuffer) agg;
            HashMap<Object, Object> partialResult = (HashMap<Object, Object>) this.internalMergeOI
                .getMap(partial);
            for (Object i : partialResult.keySet()) {
                this.putIntoSet(i, partialResult.get(i), myagg);
            }
        }

        @Override
        public void reset(AggregationBuffer buff) throws HiveException {
            MapAggBuffer arrayBuff = (MapAggBuffer) buff;
            arrayBuff.collectMap = new HashMap<>();
        }

        @Override
        public Object terminate(AggregationBuffer agg) throws HiveException {
            MapAggBuffer myagg = (MapAggBuffer) agg;
            HashMap<Object, Object> ret = new HashMap<>(myagg.collectMap);
            return ret;

        }

        private void putIntoSet(Object key, Object val, MapAggBuffer myagg) {
            Object keyCopy = ObjectInspectorUtils.copyToStandardObject(key,
                this.inputKeyOI);
            Object valCopy = ObjectInspectorUtils.copyToStandardObject(val,
                this.inputValOI);

            myagg.collectMap.put(keyCopy, valCopy);
        }

        @Override
        public Object terminatePartial(AggregationBuffer agg)
                throws HiveException {
            MapAggBuffer myagg = (MapAggBuffer) agg;
            HashMap<Object, Object> ret = new HashMap<>(myagg.collectMap);
            return ret;
        }
    }

}
